import torch

from torchvision.prototype import features, transforms as prototype_transforms
from torchvision.transforms.functional import pil_to_tensor, to_pil_image
from torchvision.utils import draw_bounding_boxes

import transforms as reference_transforms
from train import get_dataset

seed = 1234


def draw(sample):
    image, targets = sample
    annotated_image = draw_bounding_boxes(image, targets["boxes"], width=4)
    to_pil_image(annotated_image).show()


dataset, _ = get_dataset(
    name="coco",
    image_set="train",
    transform=None,
    data_path=...,
)

fill = [123.0, 117.0, 104.0]

reference_sample = (
    pil_to_tensor(dataset[0][0]),
    dict(boxes=dataset[0][1]["boxes"]),
)
prototype_sample = (
    features.Image(reference_sample[0].clone()),
    dict(
        boxes=features.BoundingBox(
            reference_sample[1]["boxes"].clone(),
            format="xyxy",
            image_size=reference_sample[0].shape[-2:],
        )
    ),
)


reference_transform = reference_transforms.RandomZoomOut(fill=fill, p=1.0)
torch.manual_seed(seed)
reference_transformed_sample = reference_transform(*reference_sample)
draw(reference_transformed_sample)

prototype_transform = prototype_transforms.RandomZoomOut(fill=fill, p=1.0)
torch.manual_seed(seed)
prototype_transformed_sample = prototype_transform(prototype_sample)
draw(prototype_transformed_sample)

torch.testing.assert_close(prototype_transformed_sample, reference_transformed_sample, atol=0, rtol=0)